'''
修正版说明（主要改动点）
1) 将滑点设置为 0（与注释一致）。可根据需要调整为 FixedSlippage(0.001)。
2) 使用 collections.deque 限制 g.slope_series 长度为 g.M，避免无限增长。
3) 统一 zscore 的标准差计算为 ddof=1（样本标准差）。
4) 修复 adjust_position 中对 context.portfolio.positions 的访问（避免 KeyError）。
5) 区分 BUY 与 KEEP：KEEP 只维持现有仓位，不主动开新仓。
6) order/成交判定更健壮，增加 try/except，避免因 order=None 抛异常。
7) 修正并同步注释（例如止损阈值为 -20 表示 20% 止损）。
8) 改进平仓/开仓判断逻辑，适配平台返回的 Order 对象可能为 None。
'''

from jqdata import *
import numpy as np
import math
from collections import deque

# 初始化函数 
def initialize(context):
    # 设定沪深300作为基准
    set_benchmark('000300.XSHG')
    # 用真实价格交易
    set_option('use_real_price', True)
    # 打开防未来函数
    set_option("avoid_future_data", True)
    # 将滑点设置为0（注：可根据实盘/回测需求调整）
    set_slippage(FixedSlippage(0.0))
    # 设置交易成本万分之三，印花税为0，最低佣金5元（类型 fund）
    set_order_cost(OrderCost(open_tax=0, close_tax=0, open_commission=0.0003, close_commission=0.0003, close_today_commission=0, min_commission=5),
                   type='fund')
    # 过滤order中低于error级别的日志
    log.set_level('order', 'error')

    # ========== 全局参数 ==========
    # 股票池（ETF）
    g.stock_pool = [
        '159915.XSHE', # 易方达创业板ETF
        '510300.XSHG', # 华泰柏瑞沪深300ETF
        '510500.XSHG', # 南方中证500ETF
    ]
    # 动量轮动参数
    g.stock_num = 1  # 买入评分最高的前 stock_num 只股票（这里为 1）
    g.momentum_day = 29  # 计算动量的天数

    # rsrs 择时参数
    g.ref_stock = '000300.XSHG'  # 用指数做择时基础
    g.N = 18   # 计算斜率使用的窗口
    g.M = 600  # 构建 zscore 的历史长度
    g.score_threshold = 0.7  # rsrs 标准分阈值

    # ma 择时参数
    g.mean_day = 20
    g.mean_diff_day = 3

    # 初始化 slope_series，使用 deque 限长（避免长期回测内存增长）
    init_slopes = initial_slope_series()  # 返回长度为 g.M 的列表
    # 去掉最新一个以避免在后续 append 时重复（与原代码逻辑一致）
    g.slope_series = deque(init_slopes[:-1], maxlen=g.M)

    # 设置交易时间，每天运行
    run_daily(my_trade, time='11:30', reference_security='000300.XSHG')
    run_daily(check_lose, time='open', reference_security='000300.XSHG')
    run_daily(print_trade_info, time='15:30', reference_security='000300.XSHG')


# ===================== 选股模块 =====================
def get_rank(stock_pool):
    """
    基于对数价格回归的动量评分：
    slope->年化收益；r2->拟合优度；score = annualized_returns * r_squared
    返回 top g.stock_num 的代码列表（降序）
    """
    score_list = []
    for stock in stock_pool:
        try:
            data = attribute_history(stock, g.momentum_day, '1d', ['close'])
            if data is None or len(data.close) < g.momentum_day:
                score_list.append(-9999)  # 数据不足，置极小值
                continue
            y = np.log(data.close.values)
            x = np.arange(y.size)
            slope, intercept = np.polyfit(x, y, 1)
            # 年化收益 (注意：slope 为日对数收益率)
            annualized_returns = math.exp(slope * 250) - 1
            # r_squared（样本方差，ddof=1）
            sse = sum((y - (slope * x + intercept)) ** 2)
            sst = (len(y) - 1) * np.var(y, ddof=1)
            r_squared = 1 - sse / sst if sst != 0 else 0
            score = annualized_returns * r_squared
            score_list.append(score)
        except Exception as e:
            # 若回归出错，记录非常低分并继续
            log.info("get_rank error for {}: {}".format(stock, str(e)))
            score_list.append(-9999)

    stock_dict = dict(zip(stock_pool, score_list))
    sort_list = sorted(stock_dict.items(), key=lambda item: item[1], reverse=True)
    code_list = [item[0] for item in sort_list]
    rank_stock = code_list[:g.stock_num]
    print("Rank list top5:", code_list[:5])
    return rank_stock


# ===================== 择时模块 =====================
def get_ols(x, y):
    """
    对 x, y 做线性回归，返回 (intercept, slope, r2)
    使用 ddof=1 计算样本方差以保持统计一致性
    """
    slope, intercept = np.polyfit(x, y, 1)
    sse = sum((y - (slope * x + intercept)) ** 2)
    sst = (len(y) - 1) * np.var(y, ddof=1)
    r2 = 1 - sse / sst if sst != 0 else 0
    return (intercept, slope, r2)


def initial_slope_series():
    """
    通过前 (N + M) 天的 high/low 做 N 天滑窗回归，生成 M 个斜率值（长度为 M）
    """
    data = attribute_history(g.ref_stock, g.N + g.M, '1d', ['high', 'low'])
    if data is None or len(data.high) < g.N + g.M:
        # 数据不足时返回一个长度为 g.M 的 0 列表（回测早期保护）
        return [0.0] * g.M
    slopes = []
    for i in range(g.M):
        window_low = data.low[i:i + g.N].values
        window_high = data.high[i:i + g.N].values
        try:
            _, slope, _ = get_ols(window_low, window_high)
            slopes.append(slope)
        except Exception:
            slopes.append(0.0)
    return slopes


def get_zscore(slope_series):
    """
    计算 zscore，使用样本标准差 ddof=1（与 get_ols 中的方差一致）
    slope_series: list-like，至少包含 1 个元素
    """
    arr = np.array(slope_series)
    if arr.size < 2:
        return 0.0
    mean = np.mean(arr)
    std = np.std(arr, ddof=1)
    if std == 0:
        return 0.0
    return (arr[-1] - mean) / std


def get_timing_signal(stock):
    """
    返回 'BUY' / 'SELL' / 'KEEP'
    逻辑：
      - MA 条件（短期 MA 相较上一段是否上升/下降）
      - RSRS 条件：rsrs_score = zscore(slope_series_lastM) * r2
    BUY 需要两者同时满足；SELL 同理；KEEP 表示不改变（此版本 KEEP 不会引发新开仓）
    """
    # MA 信号
    close_data = attribute_history(g.ref_stock, g.mean_day + g.mean_diff_day, '1d', ['close'])
    if close_data is None or len(close_data.close) < g.mean_day + g.mean_diff_day:
        # 数据不足，返回 KEEP（保守处理）
        return "KEEP"
    today_MA = close_data.close[g.mean_diff_day:].mean()
    before_MA = close_data.close[:-g.mean_diff_day].mean()

    # RSRS 信号
    high_low_data = attribute_history(g.ref_stock, g.N, '1d', ['high', 'low'])
    if high_low_data is None or len(high_low_data.high) < g.N:
        return "KEEP"
    try:
        intercept, slope, r2 = get_ols(high_low_data.low.values, high_low_data.high.values)
    except Exception:
        return "KEEP"

    # 维护 deque（已在 initialize 中创建）
    g.slope_series.append(slope)
    # 取最近 g.M 个斜率计算 zscore
    rsrs_z = get_zscore(list(g.slope_series))
    rsrs_score = rsrs_z * r2  # 修正标准分 = zscore * r2

    # 综合判断所有信号
    if rsrs_score > g.score_threshold and today_MA > before_MA:
        print('BUY (rsrs_score={:.4f}, MA_up True)'.format(rsrs_score))
        return "BUY"
    elif rsrs_score < -g.score_threshold and today_MA < before_MA:
        print('SELL (rsrs_score={:.4f}, MA_up False)'.format(rsrs_score))
        return "SELL"
    else:
        print('KEEP (rsrs_score={:.4f})'.format(rsrs_score))
        return "KEEP"


# ===================== 过滤模块 =====================
def filter_paused_stock(stock_list):
    current_data = get_current_data()
    return [stock for stock in stock_list if stock in current_data and not current_data[stock].paused]


def filter_st_stock(stock_list):
    current_data = get_current_data()
    return [stock for stock in stock_list
            if stock in current_data
            and not current_data[stock].is_st
            and 'ST' not in current_data[stock].name
            and '*' not in current_data[stock].name
            and '退' not in current_data[stock].name]


def filter_limitup_stock(context, stock_list):
    last_prices = history(1, unit='1m', field='close', security_list=stock_list)
    current_data = get_current_data()
    filtered = []
    for stock in stock_list:
        if stock in context.portfolio.positions and context.portfolio.positions[stock].total_amount > 0:
            filtered.append(stock)
            continue
        try:
            if last_prices[stock][-1] < current_data[stock].high_limit:
                filtered.append(stock)
        except Exception:
            # 若数据异常，保守地过滤掉该标的
            continue
    return filtered


def filter_limitdown_stock(context, stock_list):
    last_prices = history(1, unit='1m', field='close', security_list=stock_list)
    current_data = get_current_data()
    filtered = []
    for stock in stock_list:
        if stock in context.portfolio.positions and context.portfolio.positions[stock].total_amount > 0:
            filtered.append(stock)
            continue
        try:
            if last_prices[stock][-1] > current_data[stock].low_limit:
                filtered.append(stock)
        except Exception:
            continue
    return filtered


# ===================== 交易模块 =====================
def order_target_value_(security, value):
    """
    下单封装：返回 Order 或 None
    """
    try:
        if value == 0:
            log.debug("Selling out %s" % (security))
        else:
            log.debug("Order %s to value %f" % (security, value))
        return order_target_value(security, value)
    except Exception as e:
        log.info("order_target_value_ error for {}: {}".format(security, str(e)))
        return None


def open_position(security, value):
    """
    买入指定价值的证券，若报单创建且有成交（filled>0）则返回 True
    """
    order = order_target_value_(security, value)
    if order is None:
        return False
    # 有些平台在创建报单时不会立即返回 filled 信息，允许部分成交视为成功
    try:
        if hasattr(order, 'filled') and order.filled > 0:
            return True
        # 若没有 filled 字段或为 0，则仍认为没有成交
        return False
    except Exception:
        return False


def close_position(position):
    """
    平仓并判断是否全部成交：若报单为 None 或未全部成交则返回 False
    """
    security = position.security
    order = order_target_value_(security, 0)
    if order is None:
        return False
    try:
        # 一些平台用 status 或 filled/amount 字段，稳健检查 filled == amount
        if hasattr(order, 'filled') and hasattr(order, 'amount'):
            return order.filled == order.amount
        # 如果没有 amount 字段，但有 status，尝试容错判断（可能因平台差异）
        if hasattr(order, 'status'):
            return order.status == OrderStatus.held
    except Exception:
        pass
    # 默认返回 False（表示不能确认全部成交）
    return False


def adjust_position(context, buy_stocks, allow_open_new=True):
    """
    调仓：
      - 平掉当前持仓中不在 buy_stocks 的标的
      - 若 allow_open_new=True，按 g.stock_num 与当前持仓数量差额开仓（平均分配剩余可用现金）
    注意：buy_stocks 是候选池（已过滤过）
    """
    # 平掉不在 buy_stocks 的持仓
    current_positions = list(context.portfolio.positions.keys())
    for stock in current_positions:
        try:
            if stock not in buy_stocks:
                log.info("[%s]已不在应买入列表中，尝试平仓" % (stock))
                position = context.portfolio.positions[stock]
                close_position(position)
            else:
                log.info("[%s]已经持有无需重复买入" % (stock))
        except Exception as e:
            log.info("adjust_position check pos error: {}".format(str(e)))

    # 重新计算持仓数（只计实际持有数量 >0）
    position_count = sum(1 for p in context.portfolio.positions.values() if p.total_amount > 0)

    # 仅在允许新开仓且目标仓位大于现有仓位时开仓
    if not allow_open_new:
        return

    slots_to_open = g.stock_num - position_count
    if slots_to_open <= 0:
        return

    # 平均分配当前可用现金（注意：若有未成交的平仓单，cash 可能未释放）
    if context.portfolio.cash <= 0:
        return
    value_per_slot = context.portfolio.cash / slots_to_open

    for stock in buy_stocks:
        # 仅对未持有或持仓为0的标的开仓
        held_amt = context.portfolio.positions[stock].total_amount if stock in context.portfolio.positions else 0
        if held_amt == 0:
            if open_position(stock, value_per_slot):
                # 成功开仓后检查是否已经达到目标仓位数
                position_count = sum(1 for p in context.portfolio.positions.values() if p.total_amount > 0)
                if position_count >= g.stock_num:
                    break


def my_trade(context):
    """
    主交易函数（每日运行）
    流程：
      1) 选股并逐步过滤
      2) 获取择时信号
      3) SELL -> 全部平仓； BUY -> 调仓（开/平）； KEEP -> 仅维持现有仓位（不主动建新仓）
    """
    # 1) 选股并过滤
    try:
        check_out_list = get_rank(g.stock_pool)
        check_out_list = filter_st_stock(check_out_list)
        check_out_list = filter_limitup_stock(context, check_out_list)
        check_out_list = filter_limitdown_stock(context, check_out_list)
        check_out_list = filter_paused_stock(check_out_list)
    except Exception as e:
        log.info("选股/过滤异常: {}".format(str(e)))
        check_out_list = []

    print('今日自选股:{}'.format(check_out_list))

    # 2) 获取择时信号
    timing_signal = get_timing_signal(g.ref_stock)
    print('今日择时信号:{}'.format(timing_signal))

    # 3) 交易逻辑
    if timing_signal == 'SELL':
        # 平掉所有持仓
        for position in list(context.portfolio.positions.values()):
            try:
                close_position(position)
            except Exception as e:
                log.info("平仓异常: {}".format(str(e)))
    elif timing_signal == 'BUY':
        # 调仓并允许开新仓
        adjust_position(context, check_out_list, allow_open_new=True)
    elif timing_signal == 'KEEP':
        # 仅维护现有仓位，不主动开新仓
        adjust_position(context, [], allow_open_new=False)
    else:
        pass


# ===================== 止损模块 =====================
def check_lose(context):
    """
    止损检查：当某仓位浮动亏损 <= -20% 时触发平仓（代码与注释一致：20% 止损）
    """
    for position in list(context.portfolio.positions.values()):
        try:
            securities = position.security
            cost = position.avg_cost
            price = position.price
            if cost is None or cost == 0:
                continue
            ret = 100 * (price / cost - 1)
            value = position.value
            # 设定 20% 止损（小幅止损）
            if ret <= -20:
                order_target_value(position.security, 0)
                print("！！！！！！触发止损信号: 标的={},标的价值={},浮动盈亏={}% ！！！！！！"
                      .format(securities, format(value, '.2f'), format(ret, '.2f')))
        except Exception as e:
            log.info("check_lose error: {}".format(str(e)))


# ===================== 复盘/打印模块 =====================
def print_trade_info(context):
    # 打印当天成交记录
    try:
        trades = get_trades()
        for _trade in trades.values():
            print('成交记录：' + str(_trade))
    except Exception:
        pass

    # 打印账户信息
    for position in list(context.portfolio.positions.values()):
        try:
            securities = position.security
            cost = position.avg_cost
            price = position.price
            ret = 100 * (price / cost - 1) if cost and cost > 0 else 0
            value = position.value
            amount = position.total_amount
            print('代码:{}'.format(securities))
            print('成本价:{}'.format(format(cost, '.2f')))
            print('现价:{}'.format(price))
            print('收益率:{}%'.format(format(ret, '.2f')))
            print('持仓(股):{}'.format(amount))
            print('市值:{}'.format(format(value, '.2f')))
        except Exception:
            continue
    print('一天结束')
    print('———————————————————————————————————————分割线————————————————————————————————————————')